// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/filter_project.h>
#include <pollux/core/expressions.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/field_reference.h>

namespace kumo::pollux::exec {
    namespace {
        bool checkAddIdentityProjection(
            const core::TypedExprPtr &projection,
            const RowTypePtr &inputType,
            column_index_t outputChannel,
            std::vector<IdentityProjection> &identityProjections) {
            if (auto field = core::TypedExprs::asFieldAccess(projection)) {
                const auto &inputs = field->inputs();
                if (inputs.empty() ||
                    (inputs.size() == 1 &&
                     dynamic_cast<const core::InputTypedExpr *>(inputs[0].get()))) {
                    const auto inputChannel = inputType->getChildIdx(field->name());
                    identityProjections.emplace_back(inputChannel, outputChannel);
                    return true;
                }
            }

            return false;
        }

        // Split stats to attrbitute cardinality reduction to the Filter node.
        std::vector<OperatorStats> splitStats(
            const OperatorStats &combinedStats,
            const core::PlanNodeId &filterNodeId) {
            OperatorStats filterStats;

            filterStats.operatorId = combinedStats.operatorId;
            filterStats.pipelineId = combinedStats.pipelineId;
            filterStats.planNodeId = filterNodeId;
            filterStats.operatorType = combinedStats.operatorType;
            filterStats.numDrivers = combinedStats.numDrivers;

            filterStats.inputBytes = combinedStats.inputBytes;
            filterStats.inputPositions = combinedStats.inputPositions;
            filterStats.inputVectors = combinedStats.inputVectors;

            // Estimate Filter's output bytes based on cardinality change.
            const double filterRate = combinedStats.inputPositions > 0
                                          ? (combinedStats.outputPositions * 1.0 / combinedStats.inputPositions)
                                          : 1.0;

            filterStats.outputBytes = (uint64_t) (filterStats.inputBytes * filterRate);
            filterStats.outputPositions = combinedStats.outputPositions;
            filterStats.outputVectors = combinedStats.outputVectors;

            auto projectStats = combinedStats;
            projectStats.inputBytes = filterStats.outputBytes;
            projectStats.inputPositions = filterStats.outputPositions;
            projectStats.inputVectors = filterStats.outputVectors;

            return {std::move(projectStats), std::move(filterStats)};
        }
    } // namespace

    FilterProject::FilterProject(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::FilterNode> &filter,
        const std::shared_ptr<const core::ProjectNode> &project)
        : Operator(
              driverCtx,
              project ? project->outputType() : filter->outputType(),
              operatorId,
              project ? project->id() : filter->id(),
              "FilterProject"),
          hasFilter_(filter != nullptr),
          project_(project),
          filter_(filter) {
        if (filter_ != nullptr && project_ != nullptr) {
            stats().withWLock([&](auto &stats) {
                stats.setStatSplitter(
                    [filterId = filter_->id()](const auto &combinedStats) {
                        return splitStats(combinedStats, filterId);
                    });
            });
        }
    }

    void FilterProject::initialize() {
        Operator::initialize();
        std::vector<core::TypedExprPtr> allExprs;
        if (hasFilter_) {
            POLLUX_CHECK_NOT_NULL(filter_);
            allExprs.push_back(filter_->filter());
        }
        if (project_) {
            const auto &inputType = project_->sources()[0]->outputType();
            for (column_index_t i = 0; i < project_->projections().size(); i++) {
                auto &projection = project_->projections()[i];
                bool identityProjection = checkAddIdentityProjection(
                    projection, inputType, i, identityProjections_);
                if (!identityProjection) {
                    allExprs.push_back(projection);
                    resultProjections_.emplace_back(allExprs.size() - 1, i);
                }
            }
        } else {
            for (column_index_t i = 0; i < outputType_->size(); ++i) {
                identityProjections_.emplace_back(i, i);
            }
            isIdentityProjection_ = true;
        }
        numExprs_ = allExprs.size();
        exprs_ = makeExprSetFromFlag(std::move(allExprs), operatorCtx_->execCtx());

        if (numExprs_ > 0 && !identityProjections_.empty()) {
            const auto inputType = project_
                                       ? project_->sources()[0]->outputType()
                                       : filter_->sources()[0]->outputType();
            std::unordered_set<uint32_t> distinctFieldIndices;
            for (auto field: exprs_->distinctFields()) {
                auto fieldIndex = inputType->getChildIdx(field->name());
                distinctFieldIndices.insert(fieldIndex);
            }
            for (auto identityField: identityProjections_) {
                if (distinctFieldIndices.find(identityField.inputChannel) !=
                    distinctFieldIndices.end()) {
                    multiplyReferencedFieldIndices_.push_back(identityField.inputChannel);
                }
            }
        }
        filter_.reset();
        project_.reset();
    }

    void FilterProject::addInput(RowVectorPtr input) {
        input_ = std::move(input);
        numProcessedInputRows_ = 0;
    }

    bool FilterProject::allInputProcessed() {
        if (!input_) {
            return true;
        }
        if (numProcessedInputRows_ == input_->size()) {
            input_ = nullptr;
            return true;
        }
        return false;
    }

    bool FilterProject::isFinished() {
        return noMoreInput_ && allInputProcessed();
    }

    RowVectorPtr FilterProject::getOutput() {
        if (allInputProcessed()) {
            return nullptr;
        }

        vector_size_t size = input_->size();
        LocalSelectivityVector localRows(*operatorCtx_->execCtx(), size);
        auto *rows = localRows.get();
        POLLUX_DCHECK_NOT_NULL(rows);
        rows->setAll();
        EvalCtx evalCtx(operatorCtx_->execCtx(), exprs_.get(), input_.get());

        // Pre-load lazy vectors which are referenced by both expressions and identity
        // projections.
        for (auto fieldIdx: multiplyReferencedFieldIndices_) {
            evalCtx.ensureFieldLoaded(fieldIdx, *rows);
        }

        if (!hasFilter_) {
            numProcessedInputRows_ = size;
            POLLUX_CHECK(!isIdentityProjection_);
            auto results = project(*rows, evalCtx);

            return fillOutput(size, nullptr, results);
        }

        // evaluate filter
        auto numOut = filter(evalCtx, *rows);
        numProcessedInputRows_ = size;
        if (numOut == 0) {
            // no rows passed the filer
            input_ = nullptr;
            return nullptr;
        }

        bool allRowsSelected = (numOut == size);

        // evaluate projections (if present)
        std::vector<VectorPtr> results;
        if (!isIdentityProjection_) {
            if (!allRowsSelected) {
                rows->setFromBits(filterEvalCtx_.selectedBits->as<uint64_t>(), size);
            }
            results = project(*rows, evalCtx);
        }

        return fillOutput(
            numOut,
            allRowsSelected ? nullptr : filterEvalCtx_.selectedIndices,
            results);
    }

    std::vector<VectorPtr> FilterProject::project(
        const SelectivityVector &rows,
        EvalCtx &evalCtx) {
        std::vector<VectorPtr> results;
        exprs_->eval(
            hasFilter_ ? 1 : 0, numExprs_, !hasFilter_, rows, evalCtx, results);
        return results;
    }

    vector_size_t FilterProject::filter(
        EvalCtx &evalCtx,
        const SelectivityVector &allRows) {
        std::vector<VectorPtr> results;
        exprs_->eval(0, 1, true, allRows, evalCtx, results);
        return processFilterResults(results[0], allRows, filterEvalCtx_, pool());
    }
} // namespace kumo::pollux::exec
